#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
根据 NCBI datasets 下载生成的 dataset_catalog.json，
提取对应的基因组 FASTA (genomic.fna) 和注释文件 (gff3) 的 FTP 下载链接。

此脚本不下载数据，只生成每个 assembly 各自的 URL 列表文件：

输出目录结构示例：
    out_dir/
        GCA_012345678.1/
            GCA_012345678.1.urls
        GCA_023456789.1/
            GCA_023456789.1.urls

使用示例：
    python parse_dataset_catalog.py --catalog dataset_catalog.json --out_dir urls/
"""

import json
import re
from argparse import ArgumentParser, Namespace
from pathlib import Path


def parse_arguments() -> Namespace:
    """解析命令行参数"""
    parser = ArgumentParser(description="从 NCBI dataset_catalog.json 中提取基因组下载 URL")
    parser.add_argument("--catalog", required=True, help="dataset_catalog.json 文件路径")
    parser.add_argument("--out_dir", required=True, help="输出目录，用于存放生成的 .urls 文件")
    return parser.parse_args()


def build_ftp_prefix(accession: str) -> str:
    """
    根据 Accession 构建 FTP 路径前缀。

    例如：
        GCA_039239655.1 →
        GCA / 039 / 239 / 655 / GCA_039239655.1

    返回形如：
        GCA/039/239/655
    """
    match = re.match(r"^(GCA|GCF)_(\d{9})", accession)
    if not match:
        return ""
    prefix = match.group(1)
    digits = match.group(2)
    return f"{prefix}/{digits[:3]}/{digits[3:6]}/{digits[6:]}"


def extract_urls(catalog_path: str, out_dir: str) -> None:
    """解析 dataset_catalog.json 并生成 URL 列表文件"""
    catalog_path = Path(catalog_path)
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    with catalog_path.open("r", encoding="utf-8") as fh:
        catalog = json.load(fh)

    for asm in catalog.get("assemblies", []):
        accession = asm.get("accession")
        files = asm.get("files", [])

        if not accession or not files:
            continue

        # 定位 genomic.fna 和 gff3
        fna_entry = next((f for f in files if f.get("fileType") == "GENOMIC_NUCLEOTIDE_FASTA"), None)
        gff_entry = next((f for f in files if f.get("fileType") == "GFF3"), None)
        if not fna_entry or not gff_entry:
            continue

        ftp_prefix = build_ftp_prefix(accession)
        if not ftp_prefix:
            continue

        # 从 path 解析出 NCBI ID，例如：
        #   "/.../GCA_039239655.1_genomic.fna" → "GCA_039239655.1"
        match_id = re.search(r"/([^/]+)_genomic\.fna$", fna_entry.get("filePath", ""))
        if not match_id:
            continue
        ncbi_id = match_id.group(1)

        # 输出文件路径: out_dir/GCA_XXXXX.1/GCA_XXXXX.1.urls
        acc_dir = out_dir / accession
        acc_dir.mkdir(parents=True, exist_ok=True)
        url_file = acc_dir / f"{accession}.urls"

        with url_file.open("w", encoding="utf-8") as out:
            # genomic.fna URL
            fna_name = Path(fna_entry["filePath"]).name + ".gz"
            fna_url = f"https://ftp.ncbi.nlm.nih.gov/genomes/all/{ftp_prefix}/{ncbi_id}/{fna_name}"
            out.write(fna_url + "\n")

            # gff3 URL（名为 {ncbi_id}_{gff}.gz）
            gff_name = Path(gff_entry["filePath"]).name + ".gz"
            gff_url = f"https://ftp.ncbi.nlm.nih.gov/genomes/all/{ftp_prefix}/{ncbi_id}/{ncbi_id}_{gff_name}"
            out.write(gff_url + "\n")


def main():
    args = parse_arguments()
    extract_urls(args.catalog, args.out_dir)


if __name__ == "__main__":
    main()
